"""
Пакет содержит функции для кластеризации объектов и расчёта дистанций
"""
import pandas as pd
from sklearn.cluster import KMeans
from scipy.spatial import distance

def clustering(df_objects, count_clusters, base_point = None, max_iter = 200):
    """
    Основная функция кластеризация, вызывает последовательно этапы кластеризации

    :param df_objects: Сет данных объектов с координатами x,y,z 
    :param count_clusters: Количество кластеров
    :param base_point: Координаты базовой станции (по умолчанию None)
    :param max_iter: Максимальное Количество иттераций класстеризации (по умолчанию 200)

    :return: Сеты с данными объектов и кластеров
    """

    df_objects, df_clusters = clustering_kmeans(df_objects, count_clusters, max_iter)

    df_objects = objects_distance_to_centroids(df_objects, df_clusters)

    df_objects, df_clusters = set_leaders(df_objects, df_clusters)

    if base_point is not None:

        df_objects = objects_distance_to_base(df_objects, base_point)

    return df_objects, df_clusters

def clustering_kmeans(df_objects, count_clusters, max_iter=200):
    """
    Кластеризация объектов алгоритмом K-Means

    :param df_objects: Сет данных объектов с координатами x,y,z 
    :param count_clusters: Количество кластеров
    :param max_iter: Количество итераций алгоритма (по умолчанию 200)

    :return: Сеты данных Pandas
    """
    kmeans = run_kmeans(df_objects[['x', 'y', 'z']], count_clusters, max_iter)

    print(f"Итоговое количество кластеров: {kmeans.n_clusters}", end="\n\n")

    df_objects['cluster_id'] = kmeans.labels_

    df_clusters = pd.DataFrame(kmeans.cluster_centers_, columns=['center_x','center_y','center_z'])

    df_clusters['cluster_id'] = range(len(df_clusters))

    df_clusters = df_clusters[['cluster_id', 'center_x', 'center_y', 'center_z']]

    df_clusters['count_objects'] = df_objects['cluster_id'].value_counts()

    count_objects_list = df_clusters['count_objects'].tolist()
    count_objects_str = ', '.join(map(str, count_objects_list))

    print(f"Количество объектов в кластерах: {count_objects_str}", end="\n\n")

    return df_objects, df_clusters

def run_kmeans(df_objects, n_clusters, random_state=0, max_iter=200):
    """
    Запуск кластеризации в цикле для постепенного уменьшения количества кластеров.

    :param df_objects: Сет данных объектов с координатами x,y,z 
    :param n_clusters: Количество кластеров
    :param random_state: Определяет генерацию случайных чисел для инициализации центроида. (по умолчанию 0)
    :param max_iter: Количество итераций алгоритма (по умолчанию 200)

    :return: Результат kmeans
    """
    try:
        kmeans = KMeans(n_clusters=n_clusters, random_state=random_state, n_init="auto", max_iter=max_iter)
        kmeans.fit(df_objects)
        cluster_sizes = pd.Series(kmeans.labels_).value_counts()
        if (cluster_sizes < 3).any():
            print(f"Ошибка кластеризации на {n_clusters} кластеров. Попытка кластеризации на {n_clusters - 1} кластеров")
            return run_kmeans(df_objects, n_clusters - 1)
        else:
            return kmeans
    except ValueError:
        print(f"Ошибка при количестве {n_clusters} кластеров есть кластеры с менее 3 объектами. Выполняем кластеризацию на {n_clusters - 1} кластеров")
        return run_kmeans(df_objects, n_clusters - 1)

def get_by_field(df, field, value):
    """
    Вспомогательная функция для выбора данных в датасете по значению поля

    :param df: Дата сет
    :param field: Поле
    :param value: Значение для поиска

    :return: Значение
    """
    return df[df[field] == value]

def objects_distance_to_centroids(df_objects, df_clusters):
    """
    Вычисляет расстояние от объектов до центроидов кластеров.

    :param df_objects: Сет данных с координатами объектов.
    :param df_clusters: Сет данных с информацией о кластерах.

    :return: Сет данных с добавленным столбцом 'distance_to_centroid'.
    """
    for i, row in df_objects.iterrows():
        coordinates_array = [row['x'], row['y'], row['z']]
        cluster_data = df_clusters[df_clusters['cluster_id'] == row['cluster_id']]
        coordinates_centroid = [cluster_data.iloc[0]['center_x'], cluster_data.iloc[0]['center_y'], cluster_data.iloc[0]['center_z']]
        distance_to_centroid = distance.euclidean(coordinates_array, coordinates_centroid)
        df_objects.at[i, 'distance_to_centroid'] = distance_to_centroid

    return df_objects

def objects_distance_to_base(df_objects, base):
    """
    Вычисляет расстояние от объектов до базовой станции.

    :param df_objects: Сет данных с координатами объектов.
    :param base: Координаты базовой станции в формате [x, y, z].

    :return: Сет данных с добавленным столбцом 'distance_to_base'.
    """
    df_base = pd.DataFrame([base] * len(df_objects), columns=['x', 'y', 'z'])

    df_objects['distance_to_base'] = distance.cdist(df_objects[['x', 'y', 'z']], df_base[['x', 'y', 'z']]).min(axis=1)

    return df_objects

def objects_distance_to_leaders(df_objects):
    """
    Вычисляет расстояние от объектов до главных в кластере (leader_x, leader_y, leader_z).

    :param df_objects: Сет данных с координатами объектов.

    :return: Сет данных с добавленным столбцом 'distance_to_leader'.
    """
    df_objects['distance_to_leader'] = distance.cdist(df_objects[['x','y','z']], df_objects[['leader_x', 'leader_y', 'leader_z']]).min(axis=1)

    return df_objects

def set_base_connectors(df_objects, df_clusters, base_connections):
    """
    Устанавливает коннекторы с базовой станцией для объектов и кластеров.

    :param df_objects: Сет данных с информацией об объектах.
    :param df_clusters: Сет данных с информацией о кластерах.
    :param base_connections: Количество ближайших лидеров для базовых коннекторов.

    :return: Модифицированные сеты данных df_objects и df_clusters.
    """
    df_leaders = df_objects[df_objects['leader']]

    df_nearest_leaders = df_leaders.nsmallest(base_connections, 'distance_to_base')

    df_objects['base_connector'] = df_objects.index.isin(df_nearest_leaders.index)

    df_clusters['base_connector'] = df_clusters['leader_id'].isin(df_nearest_leaders['object_id'])

    distance_dict = df_leaders.set_index('object_id')['distance_to_base'].to_dict()

    df_clusters['distance_to_base'] = df_clusters['leader_id'].map(distance_dict)

    return df_objects, df_clusters

def get_base_connectors(df_objects, columns_to_return=None):
    """
    Возвращает коннекторов к базовой станции (объекты, ближайшие к базовой станции).

    :param df_objects: Сет данных с информацией об объектах.
    :param columns_to_return: Список столбцов для возврата (по умолчанию None).

    :return: Сет данных с базовыми коннекторами.
    """
    base_connectors = df_objects[df_objects['base_connector']]

    if columns_to_return is not None:
        base_connectors = df_objects[columns_to_return]

    return base_connectors

def get_leaders(df_objects, columns_to_return=None):
    """
    Возвращает информацию о лидерах (объекты, являющиеся лидерами).

    :param df_objects: Сет данных с информацией об объектах.
    :param columns_to_return: Список столбцов для возврата (по умолчанию None).

    :return: Сет данных с информацией о лидерах.
    """
    df_leaders = df_objects[df_objects['leader']]

    if columns_to_return is not None:
        df_leaders = df_leaders[columns_to_return]

    return df_leaders

def set_leaders(df_objects, df_clusters):
    """
    Устанавливает лидеров для объектов и кластеров.

    :param df_objects: Сет данных с информацией об объектах.
    :param df_clusters: Сет данных с информацией о кластерах.

    :return: Модифицированные сеты данных df_objects и df_clusters.
    """
    min_distances = df_objects.groupby('cluster_id')['distance_to_centroid'].min()

    df_objects['leader'] = df_objects.apply(lambda row: row['distance_to_centroid'] == min_distances[row['cluster_id']], axis=1)

    clusters = df_objects['cluster_id'].unique()

    for _, cluster_id in enumerate(clusters):
        leader_coords = df_objects.loc[(df_objects['leader']) & (df_objects['cluster_id'] == cluster_id), ['object_id','x', 'y', 'z']]

        df_clusters.loc[df_clusters['cluster_id'] == cluster_id, ['leader_id', 'leader_x', 'leader_y', 'leader_z']] = leader_coords.values
        df_objects.loc[df_objects['cluster_id'] == cluster_id, ['leader_id', 'leader_x', 'leader_y', 'leader_z']] = leader_coords.values

    df_clusters['leader_id'] = df_clusters['leader_id'].round().astype(int)
    df_objects['leader_id'] = df_objects['leader_id'].round().astype(int)

    df_objects = objects_distance_to_leaders(df_objects)

    return df_objects, df_clusters

def print_stat(df_objects):
    """
    Выводит статистику по объектам, лидерам и кластерам.

    :param df_objects: Сет данных с информацией об объектах.

    :return: None
    """

    print(f"Количество объектов: {len(df_objects)}", end="\n\n")

    if 'shortest_path_length' in df_objects.columns:

        distance_to_base_stats = df_objects['distance_to_base'].agg(['min', 'mean', 'median', 'max'])

        shortest_path_length_stats = df_objects['shortest_path_length'].agg(['min', 'mean', 'median', 'max'])

        not_leaders = df_objects.loc[~df_objects['leader']]
        distance_to_leader_stats = not_leaders['distance_to_leader'].agg(['min', 'mean', 'median', 'max'])

        shortest_path_length_steps = not_leaders['shortest_path_steps'].agg(['min', 'mean', 'median', 'max'])

        stat = pd.concat([distance_to_base_stats, shortest_path_length_stats, distance_to_leader_stats, shortest_path_length_steps], axis=1)
        stat.index = ['min', 'mean', 'median', 'max']

        stat.rename(columns={'distance_to_base': 'Расстояние до базовой станции',
                             'shortest_path_length': 'Маршрут до базовой станции',
                             'distance_to_leader': 'Расстояние до лидера',
                             'shortest_path_steps': 'Шагов в маршруте'}, 
                             inplace=True)

        print("Данные по дистанциям:")
        print(stat, end="\n\n")

        average_steps = not_leaders['shortest_path_steps'].mean()
        print(f"Среднее количество шагов в маршрутах связи: {average_steps:.2f}", end="\n\n")

        steps_distribution = not_leaders['shortest_path_steps'].value_counts()
        # Создание DataFrame с двумя колонками
        steps_df = pd.DataFrame({'Количество шагов': steps_distribution.index, 'Количество объектов': steps_distribution.values})
        print("Распределение по количеству шагов в маршруте:", end="\n\n")
        print(steps_df, end="\n\n")

        leaders = df_objects[df_objects['leader']]

        usage_stats = leaders[['object_id', 'cluster_id', 'shortest_path_usage_count']].copy()

        usage_stats.rename(columns={'object_id': 'ID лидера','cluster_id': 'ID кластера', 'shortest_path_usage_count': 'Количество использований'}, inplace=True)

        sorted_usage_stats = usage_stats.sort_values(by='Количество использований', ascending=False)

        print(sorted_usage_stats, end="\n\n")

    else:

        not_leaders = df_objects.loc[~df_objects['leader']]

        distance_to_base_stats = not_leaders['distance_to_base'].agg(['min', 'mean', 'median', 'max'])

        distance_to_centroid = not_leaders['distance_to_centroid'].agg(['min', 'mean', 'median', 'max'])

        distance_to_leader_stats = not_leaders['distance_to_leader'].agg(['min', 'mean', 'median', 'max'])

        leaders = df_objects[df_objects['leader']]

        distance_leaders_to_base = leaders['distance_to_base'].agg(['min', 'mean', 'median', 'max'])
        distance_leaders_to_base.rename('leaders_distance_to_base', inplace=True)

        connections = []
        for i, leader_point1 in enumerate(leaders.itertuples()):
            for j, leader_point2 in enumerate(leaders.itertuples()):
                if i != j:
                     # Получите координаты для текущих лидеров
                    leader_coordinates_array1 = [leader_point1.x, leader_point1.y, leader_point1.z]
                    leader_coordinates_array2 = [leader_point2.x, leader_point2.y, leader_point2.z]

                    # Вычислите евклидово расстояние между координатами
                    _distance = distance.euclidean(leader_coordinates_array1, leader_coordinates_array2)

                    # Добавьте информацию о расстоянии в список connections
                    connections.append((leader_point1.object_id, leader_point2.object_id, _distance))

        # Создайте DataFrame из списка connections
        df_connections = pd.DataFrame(connections, columns=['leader1', 'leader2', 'distance'])

        leaders_connections = df_connections['distance'].agg(['min', 'mean', 'median', 'max'])

        print(f"Количество связей между {len(leaders)} лидерами: {len(connections)}", end="\n\n")

        combined_df = pd.concat([distance_to_base_stats, distance_to_centroid, distance_to_leader_stats, distance_leaders_to_base, leaders_connections], axis=1)

        combined_df.rename(columns={'distance_to_base': 'От объекта до базовой станции',
                                    'distance_to_centroid': 'От объекта до центроида',
                                    'distance_to_leader': 'От объекта до лидера',
                                    'leaders_distance_to_base': 'От лидера до базовой станции',
                                    'distance': 'Между лидерами'}, 
                                    inplace=True)

        combined_df = combined_df.transpose()

        print("Данные по дистанциям:")
        print(combined_df)
